#include <Windows.h>
#include <WindowsX.h>
#include <tchar.h>
#include "resource.h"
#include <psapi.h>
#include <TlHelp32.h>
#pragma comment(lib, "Psapi.lib")

void AddTextToEdit(HWND hEdit, TCHAR* szFormat, ...)
{
	va_list vaList;
	va_start(vaList, szFormat);
	TCHAR szTmp[1024 * 400];
	Edit_GetText(hEdit, szTmp, _countof(szTmp));
	_vstprintf(_tcschr(szTmp, _T('\0')), szFormat, vaList);
	Edit_SetText(hEdit, szTmp);
	va_end(vaList);
	return;
}

void UpdateProcessList(HWND hCmbProcess)
{
	DWORD dwProcessIDs[1024];
	DWORD dwCount = 0;
	EnumProcesses(dwProcessIDs, 1024, &dwCount);

	for (int ii = ComboBox_GetCount(hCmbProcess) - 1; ii >= 0; ii--)
	{
		ComboBox_DeleteString(hCmbProcess, ii);
	}
	
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (INVALID_HANDLE_VALUE != hSnapshot)
	{
		PROCESSENTRY32 pe = {sizeof(PROCESSENTRY32)};
		if (Process32First(hSnapshot, &pe))
		{
			do {
				UINT nIndex = ComboBox_AddString(hCmbProcess, pe.szExeFile);
				ComboBox_SetItemData(hCmbProcess, nIndex, pe.th32ProcessID);
			} while (Process32Next(hSnapshot, &pe));
		}		
		ComboBox_SetCurSel(hCmbProcess, 0);
		CloseHandle(hSnapshot);
	}

	return;
}

int CALLBACK DlgProc(HWND hDlg, UINT nMsg, WPARAM wParam, LPARAM lParam)
{
	static HWND hCmbProcess = 0;
	switch (nMsg)
	{
		case WM_INITDIALOG:
		{
			SetWindowText(hDlg, _T("Demo0082"));

			UINT nWidth		= GetSystemMetrics(SM_CXSCREEN);
			UINT nHeight	= GetSystemMetrics(SM_CYSCREEN);
			RECT rtClient;
			GetClientRect(hDlg, &rtClient);
			SetWindowPos(hDlg, NULL, (nWidth - (rtClient.right - rtClient.left)) / 2, (nHeight - (rtClient.bottom - rtClient.top)) / 2, 0, 0, SWP_NOSIZE);
			
			hCmbProcess = GetDlgItem(hDlg, IDC_CMBPROCESS);
			UpdateProcessList(hCmbProcess);
			SendMessage(hDlg, WM_COMMAND, MAKEWPARAM(IDC_CMBPROCESS, CBN_SELCHANGE), 0);

			break;
		}
		case WM_COMMAND:
		{
			if (HIWORD(wParam) == CBN_SELCHANGE)
			{
				DWORD th32ProcessID = ComboBox_GetItemData(hCmbProcess, ComboBox_GetCurSel(hCmbProcess));
				HANDLE hSnapHeap = CreateToolhelp32Snapshot(TH32CS_SNAPHEAPLIST, th32ProcessID);
				if (INVALID_HANDLE_VALUE != hSnapHeap)
				{
					SetWindowText(GetDlgItem(hDlg, IDC_EDTINFO), _T(""));
					HEAPLIST32 hl32 = {sizeof(HEAPLIST32)};
					if (Heap32ListFirst(hSnapHeap, &hl32))
					{
						AddTextToEdit(GetDlgItem(hDlg, IDC_EDTINFO),
								      TEXT("%5s     %-10s   %-10s   %-10s   %8s\r\n"),
								      _T("ID"), _T("Flag"), _T("Handle"), _T("Address"), _T("Size"));
						
						AddTextToEdit(GetDlgItem(hDlg, IDC_EDTINFO),
									  TEXT("=============================================================\r\n"));
						DWORD nCount = 0;
						static TCHAR szInfo[1024 * 300];					
						do {
							HEAPENTRY32 he32 = {sizeof(HEAPENTRY32)};
							if (th32ProcessID == hl32.th32ProcessID && Heap32First(&he32, th32ProcessID, hl32.th32HeapID))
							{
								do {
									if (th32ProcessID == he32.th32ProcessID)
									{
										_stprintf(_tcschr(szInfo, _T('\0')),
									              TEXT("%5d     %-10s   0x%06X     0x%06X     %8u\r\n"),
												  nCount++,
												  (he32.dwFlags == LF32_FIXED) ? _T("Fixed") : (he32.dwFlags == LF32_FREE) ? _T("Free") : _T("Moveable"), 
												  he32.hHandle, 	
												  he32.dwAddress,
												  he32.dwBlockSize);
									}
								} while (Heap32Next(&he32));
							}
						} while (Heap32ListNext(hSnapHeap, &hl32));
						AddTextToEdit(GetDlgItem(hDlg, IDC_EDTINFO), szInfo);
					}
					CloseHandle(hSnapHeap);
				}
			}

			break;
		}
	
		case WM_CLOSE:
		{
			PostQuitMessage(0);
			break;
		}
	}

	return 0;
}

BOOL WINAPI _tWinMain(HINSTANCE hInstance, HINSTANCE hPreInstance, LPTSTR szCmdLine, int nShowCmd)
{
	DialogBox(hInstance, MAKEINTRESOURCE(IDD_DIALOG), NULL, DlgProc);

	return TRUE;
}